import re
from py2neo import Graph

# 步骤A: 连接到图数据库
try:
    graph = Graph("bolt://localhost:7687", auth=("neo4j", "053116wj"))
    print("成功连接到图数据库！")
except Exception as e:
    print(f"连接数据库失败：{e}")
    graph = None


# --- 核心函数：只负责执行查询并返回原始数据 ---
def execute_query_and_get_data(query: str) -> list:
    """
    执行 Cypher 查询并返回原始数据列表。
    如果发生错误，返回空列表。
    """
    if not graph:
        return []
    try:
        return graph.run(query).data()
    except Exception as e:
        print(f"查询出错: {e}")
        return []


# --- 专用的格式化函数 ---
def format_general_answer(data: list, entity_name: str, relationship_type: str) -> str:
    """格式化通用查询的回答。"""
    if not data:
        return f"抱歉，没有找到关于“{entity_name}”的{relationship_type}信息。"

    result_list = [item['name'] for item in data if 'name' in item and item['name']]

    if result_list:
        result_str = "、".join(result_list)
        return f"关于“{entity_name}”的{relationship_type}有：{result_str}。"
    else:
        return f"抱歉，没有找到相关信息。"


def format_characteristics_answer(data: list, entity_name: str) -> str:
    """
    格式化特点查询的回答，并明确区分优点和缺点。
    假设查询返回 'advantages' 和 'disadvantages' 字段。
    """
    if not data:
        return f"抱歉，没有找到关于“{entity_name}”的特性信息。"

    # 由于查询返回的是一个包含优缺点属性的节点，我们取第一个
    features = data[0]

    response_parts = []

    # 提取优点和缺点，如果不存在则设为空字符串
    characteristics = features.get('characteristics', '')
    disadvantages = features.get('disadvantages', '')

    response_parts = []

    # 检查优点是否为字符串，并用逗号分隔成列表
    if characteristics:
        characteristics_list = characteristics.split('，')  # 注意：这里用中文逗号
        characteristics_str = "、".join(characteristics_list)
        response_parts.append(f"特点为：{characteristics_str}")

    # 检查缺点是否为字符串，并用逗号分隔成列表
    if disadvantages:
        disadvantages_list = disadvantages.split('，')
        disadvantages_str = "、".join(disadvantages_list)
        response_parts.append(f"缺点有：{disadvantages_str}")

    if response_parts:
        return f"关于“{entity_name}”的特性：{'。'.join(response_parts)}。"
    else:
        return f"抱歉，没有找到关于“{entity_name}”的特性信息。"


def format_add_mount_answer(data: list, entity_name: str) -> str:
    """
    格式化添加量及感受查询的回答。
    """
    if not data:
        return f"抱歉，没有找到关于“{entity_name}”的添加量及感受信息。"

    # 假设查询返回的是一个包含 add_mount 和 evaluate 的字段
    info = data[0]
    add_mount = info.get('add_mount', '未提及')
    evaluate = info.get('evaluate', '未提及')

    return f"“{entity_name}”的添加量为：{add_mount}，口感为：{evaluate}。"

def format_add_method_answer(data: list, entity_name: str, relationship_type: str) -> str:
    """格式化添加方式的回答。"""
    if not data:
        return f"抱歉，没有找到关于“{entity_name}”的{relationship_type}信息。"

    result_list = [item['name'] for item in data if 'name' in item and item['name']]

    if result_list:
        result_str = "、".join(result_list)
        return f"“{entity_name}”的{relationship_type}为：{result_str}。"
    else:
        return f"抱歉，没有找到相关信息。"

def format_relative_sweetness_answer(data: list, entity_name: str, relationship_type: str) -> str:
    """格式化添加方式的回答。"""
    if not data:
        return f"抱歉，没有找到关于“{entity_name}”的{relationship_type}信息。"

    result_list = [item['name'] for item in data if 'name' in item and item['name']]

    if result_list:
        result_str = "、".join(result_list)
        return f"“{entity_name}”的{relationship_type}为：{result_str}。"
    else:
        return f"抱歉，没有找到相关信息。"

def format_duration_answer(data: list, entity_name: str, relationship_type: str) -> str:
    """格式化甜味持续性的回答。"""
    if not data:
        return f"抱歉，没有找到关于“{entity_name}”的{relationship_type}信息。"

    result_list = [item['name'] for item in data if 'name' in item and item['name']]

    if result_list:
        result_str = "、".join(result_list)
        return f"“{entity_name}”的{relationship_type}描述为：{result_str}。"
    else:
        return f"抱歉，没有找到相关信息。"

def format_transfer_rate_answer(data: list, entity_name: str) -> str:
    """
    格式化甜味物质烟叶烟丝中含量及迁移率的回答。
    """
    if not data:
        return f"抱歉，没有找到关于“{entity_name}”的添加量及感受信息。"

    # 假设查询返回的是一个包含 add_mount 和 evaluate 的字段
    info = data[0]
    smoke_measure_ug_per_cig = info.get('smoke_measure_ug_per_cig', '未提及')
    tobacco_measure_ug_per_cig = info.get('tobacco_measure_ug_per_cig', '未提及')
    tobacco_transfer_rate_pct = info.get('tobacco_transfer_rate_pct', '未提及')

    return f"“{entity_name}”烟气中的含量为：{smoke_measure_ug_per_cig},烟叶中含量为{tobacco_measure_ug_per_cig},烟叶烟丝中的迁移率为：{tobacco_transfer_rate_pct}。"


# --- 主问答函数：负责问题解析和调用格式化函数 ---
def answer_question(question: str) -> str:
    """
    处理用户的自然语言问题，并返回答案。
    """
    if not graph:
        return "抱歉，无法连接到图数据库。"

    # 规则1: 甜味剂类型关联的文献查询
    # 调整后的新规则1
    match = re.search(r'(?:《|“|)(.+?)(?:》|”|)(?:有什么|哪些)?(?:文献|论文)？', question)
    if match:
        sweetener_add_name = match.group(1)
        query = f"MATCH (:甜味剂类型 {{add_name: '{sweetener_add_name}'}})-[:研究论文]->(m:文献) RETURN m.paper_name AS name limit 5"
        result_data = execute_query_and_get_data(query)
        return format_general_answer(result_data, sweetener_add_name, "相关文献")

    # 规则2: 文献中提到的致甜成分查询
    match = re.search(r'(?:《|“)(.+?)(?:》|”)(提到了|涉及到)(哪些|什么)(致甜成分|致甜物质)？', question)
    if match:
        document_paper_name = match.group(1)
        query = f"MATCH (:文献 {{paper_name: '{document_paper_name}'}}) -[:致甜物质]->(e:致甜成分) RETURN e.substance_name AS name"
        result_data = execute_query_and_get_data(query)
        return format_general_answer(result_data, document_paper_name, "致甜成分")

    # 规则3：致甜成分的特点查询
    match = re.search(r'(?:《|“|)(.+?)(?:》|”|)有什么(特点|优点|特性)？', question)
    if match:
        ingredient_sub_name = match.group(1)

        # 优化后的 Cypher 查询，返回优缺点两个独立的字段
        query = f"""
        MATCH (:致甜成分 {{substance_name: '{ingredient_sub_name}'}})-[:特性]->(f:致甜成分特点) 
        RETURN f.characteristics AS characteristics, f.disadvantages AS disadvantages
        """
        result_data = execute_query_and_get_data(query)
        # 调用专门的格式化函数，并传递原始数据
        return format_characteristics_answer(result_data, ingredient_sub_name)

    # 规则4: 致甜成分应用查询
    match = re.search(r'(?:《|“|)(.+?)(?:》|”|)用在哪里？|应用是什么？', question)
    if match:
        ingredient_sub_name = match.group(1)
        query = f"MATCH (:致甜成分 {{substance_name: '{ingredient_sub_name}'}})-[:应用]->(a:致甜成分应用) RETURN a.application_brand_scene AS name"
        result_data = execute_query_and_get_data(query)
        return format_general_answer(result_data, ingredient_sub_name, "应用")

    # 规则5：查询致甜成分的添加量及感受
    match = re.search(r'(?:《|“|)(.+?)(?:》|”|)的添加量及(感受|口感)是怎样的？', question)
    if match:
        ingredient_sub_name = match.group(1)
        # 注意：这里的查询需要返回多个字段
        query = f"MATCH (:致甜成分 {{substance_name: '{ingredient_sub_name}'}})-[:添加量及感受]->(m:致甜成分添加量) RETURN m.add_mount AS add_mount, m.evaluate AS evaluate"
        result_data = execute_query_and_get_data(query)
        return format_add_mount_answer(result_data, ingredient_sub_name)

    # 规则6A: 问法一 “X的添加方式是怎样的？”
    match_method = re.search(r'(?:《|“|)(.+?)(?:》|”|)的添加(?:方式|方法)是怎样的？', question)
    if match_method:
        ingredient_sub_name = match_method.group(1)
        query = f"MATCH (:致甜成分 {{substance_name: '{ingredient_sub_name}'}})-[:添加方式]->(m:致甜成分添加方法) RETURN m.add_method AS name"
        result_data = execute_query_and_get_data(query)
        return format_add_method_answer(result_data, ingredient_sub_name, "添加方式")

    # 规则6B: 问法二 “X怎么添加？”
    match_how_to_add = re.search(r'(?:《|“|)(.+?)(?:》|”|)(怎么|如何)添加？', question)
    if match_how_to_add:
        ingredient_sub_name = match_how_to_add.group(1)
        query = f"MATCH (:致甜成分 {{substance_name: '{ingredient_sub_name}'}})-[:添加方式]->(m:致甜成分添加方法) RETURN m.add_method AS name"
        result_data = execute_query_and_get_data(query)
        return format_add_method_answer(result_data, ingredient_sub_name, "添加方式")

    # 规则7: 查询相对甜度”
    match_how_to_add = re.search(r'(?:《|“|)(.+?)(?:》|”|)的(相对甜度|甜度)是多少？', question)
    if match_how_to_add:
        ingredient_sub_name = match_how_to_add.group(1)
        query = f"MATCH (:致甜成分 {{substance_name: '{ingredient_sub_name}'}})-[:相对甜度]->(m:致甜成分相对甜度) RETURN m.relative_sweetness AS name"
        result_data = execute_query_and_get_data(query)
        return format_relative_sweetness_answer(result_data, ingredient_sub_name, "相对甜度")

    #规则8：查询甜味持久性
    match_how_to_add = re.search(r'(?:《|“|)(.+?)(?:》|”|)的甜味(?:持久性|持续性)(如何|怎么样)？', question)
    if match_how_to_add:
        ingredient_sub_name = match_how_to_add.group(1)
        query = f"MATCH (:致甜成分 {{substance_name: '{ingredient_sub_name}'}})-[:持久性]->(m:致甜成分甜味持久性) RETURN m.duration AS name"
        result_data = execute_query_and_get_data(query)
        return format_duration_answer(result_data, ingredient_sub_name, "持久性")

    #规则9：查询烟叶烟丝中甜味物质的含量及迁移率
    # 简化后的正则表达式，涵盖了各种问法
    match_how_to_add = re.search(
        r'(?:《|“|)(.+?)(?:》|”|)在(?:烟叶|烟丝|烟草)中的(?:含量|迁移率|含量及迁移率|含量和迁移率)\??', question)
    if match_how_to_add:
        ingredient_sub_name = match_how_to_add.group(1)
        query = f"MATCH (:致甜成分 {{substance_name: '{ingredient_sub_name}'}})-[:含量及迁移率]->(m:烟叶烟丝中的含量及迁移率) RETURN m.smoke_measure_ug_per_cig AS smoke_measure_ug_per_cig,m.tobacco_measure_ug_per_cig AS tobacco_measure_ug_per_cig,m.tobacco_transfer_rate_pct AS tobacco_transfer_rate_pct"
        result_data = execute_query_and_get_data(query)
        return format_transfer_rate_answer(result_data, ingredient_sub_name)

    #规则10：查询某个致甜成分在哪些文献中被提及
    # 修复后的规则10正则表达式
    match_how_to_add = re.search(r'(?:《|“|)(.+?)(?:》|”|)这个成分在哪些文献中(?:被)?提到了？', question)
    if match_how_to_add:
        ingredient_sub_name = match_how_to_add.group(1)
        query = f"MATCH (:致甜成分 {{substance_name: '{ingredient_sub_name}'}})<-[:致甜物质]-(doc:文献) RETURN doc.paper_name AS name"
        result_data = execute_query_and_get_data(query)
        return format_general_answer(result_data, ingredient_sub_name, relationship_type="相关文献")


    # 如果所有规则都不匹配
    return "抱歉，我无法理解你的问题。请尝试换一种说法，例如“天然植物提取物有什么文献？”"


# 步骤E: 测试你的问答系统
if __name__ == "__main__":
    print(">>> 测试问题1：天然植物提取物有什么文献？")
    print(answer_question("天然植物提取物有什么文献？"))
    print("\n>>> 测试问题2：《甜菊糖苷应用研究现状》提到了哪些致甜成分？")
    print(answer_question("《甜菊糖苷应用研究现状》提到了哪些致甜成分？"))
    print("\n>>> 测试问题3：甜菊糖苷（瑞鲍迪甙A）有什么特点？")
    print(answer_question("甜菊糖苷（瑞鲍迪甙A）有什么特点？"))
    print("\n>>> 测试问题4：甜菊糖苷（瑞鲍迪甙A）用在哪里？")
    print(answer_question("甜菊糖苷（瑞鲍迪甙A）用在哪里？"))
    print("\n>>> 测试问题5：甜菊糖苷（瑞鲍迪甙A）的添加量及感受是怎样的？")
    print(answer_question("甜菊糖苷（瑞鲍迪甙A）的添加量及感受是怎样的？"))
    print("\n>>> 测试问题6：大果木姜子茎挥发油怎么添加？")
    print(answer_question("大果木姜子茎挥发油怎么添加？"))
    print("\n>>> 测试问题7：甜菊糖苷（瑞鲍迪甙A）的相对甜度是多少？")
    print(answer_question("甜菊糖苷（瑞鲍迪甙A）的相对甜度是多少？"))
    print("\n>>> 测试问题8：三氯蔗糖的甜味持续性如何？")
    print(answer_question("三氯蔗糖的甜味持续性如何？"))
    print("\n>>> 测试问题9：山梨醇在烟叶中的含量及迁移率？")
    print(answer_question("山梨醇在烟叶中的含量及迁移率？"))
    print("\n>>> 测试问题10：肌醇这个成分在哪些文献中提到了？")
    print(answer_question("肌醇这个成分在哪些文献中提到了？"))
